import os
import shutil
import SimpleITK as sitk
import numpy as np
import time
import math
import multiprocessing as mp

from mytest.mycode import find_long_axis, calc_diff, antialias, geoutils, fileutils, post_process


# set the __is_pred to process the prediction segmentations
__is_pred = True

# ===================== top-level root paths for following settings =========================
data_root = 'mydata_pred' # switch between mydata_gt and mydata_pred.
image_root = 'images'
gt_root = 'groundtruth'

# ===================== basic roots for the computing ========================================
axis_calc_root = os.path.join(data_root, 'axis_calc')
rotate_root = os.path.join(data_root, 'rotate')
pred_recover_root = os.path.join(data_root, 'pred_recover')

# ============= path details for the subprocessing procedures ===============
# -------- root paths for calculating the metrics measurements of long axis locating ------------
manual_root = [
    os.path.join(axis_calc_root, 'manual_axis', 'long_axis_1'),
    os.path.join(axis_calc_root, 'manual_axis', 'new_axis_A'),
    os.path.join(axis_calc_root, 'manual_axis', 'new_axis_B')
]
metrics_root = [
    os.path.join(axis_calc_root, 'metrics', '1'),
    os.path.join(axis_calc_root, 'metrics', '2'),
    os.path.join(axis_calc_root, 'metrics', '3'),
]

# --------- root paths for calculating the long axis --------------------
output_axis_root = os.path.join(axis_calc_root, 'output_with_axis')
generated_root = os.path.join(axis_calc_root, 'generated_axis')
merged_root = os.path.join(axis_calc_root, 'merged_axis')
anti_root = os.path.join(axis_calc_root, 'results_anti')

# --------- root paths for recovering the predicted segments ------------
resampled_root = os.path.join(pred_recover_root, 'results_resampled')
recovered_root = os.path.join(pred_recover_root, 'results_recovered')
aliased_root = os.path.join(pred_recover_root, 'results_recovered')
templateinfo_root = os.path.join(pred_recover_root, 'template_info')

# --------- root paths for rotating the images to SAX view ---------------
shortAxis_root = os.path.join(rotate_root, 'output_shortAxis')
tempinfo_root = os.path.join(rotate_root, 'tempinfo')
shortAxisSeg_root = os.path.join(rotate_root, 'shortAxis_seg')
rotateBack_root = os.path.join(rotate_root, 'rotate_back')



PRED_PREFIX = 'prediction_'
NIFTI_SUFFIX = '.nii.gz'
OUTPUT_AXIS_PREFIX = 'outputAxis_'
OUTPUT_ROTATED_PREFIX = 'outputRotated_'
GENERATED_PREFIX = 'generated_'
TXT_SUFFIX = '.txt'
METRICS_PREFIX = 'metrics_'
STATISTICS_PREFIX = 'statistics_'
M0_SUFFIX = '-Segmentation.seg'
M1_SUFFIX = 'A.seg'
M2_SUFFIX = 'B.seg'
NRRD_SUFFIX = '.nrrd'
MERGED_PREFIX = 'allAxis_'
SHORTAXIS_PREFIX = 'shortAxis_'
TEMPLATEINFO_PREFIX = 'template_'
IMAGEDIR_PREFIX = 'case_'
IMAGE_NAME = 'imaging.nii.gz'
SEG_NAME = 'segmentation.nii.gz'
TEMPINFO_PREFIX = 'tempinfo_'
SEGMENTATION_PREFIX = 'segmentation_'


def calc_metrics(seq_list, man_idx):
    manuals = os.listdir(manual_root[man_idx])
    manual_dict = dict({})
    for manual_file in manuals:
        seq_number = manual_file.split('.')[0].split('-')[0]
        seq_number = ''.join(list(filter(str.isdigit, seq_number)))
        manual_dict[int(seq_number)] = manual_file

    # gen_axises = os.listdir(generated_root)
    for seq_number in seq_list:
        gen_axis_path = os.path.join(generated_root, GENERATED_PREFIX + seq_number + TXT_SUFFIX)
        
        manual_axis_path = os.path.join(manual_root[man_idx], manual_dict[int(seq_number)])
        metrics_path = os.path.join(metrics_root[man_idx], METRICS_PREFIX + seq_number + TXT_SUFFIX)
        try:
            calc_diff.calc_diff(manual_axis_path, my_axis_path=gen_axis_path, output_path=metrics_path)
        except:
            print('wrong: ', seq_number, ', manual: ', man_idx)


def calc_statistics(seq_list, man_idx):
    dist_list = []
    alpha_diff_list = []
    beta_diff_list = []
    theta_list = []
    for seq_number in seq_list:
        metrics_path = os.path.join(metrics_root[man_idx], METRICS_PREFIX + seq_number + TXT_SUFFIX)
        distance, (p_a, p_b), (m_a, m_b), theta = fileutils.read_metrics(metrics_path)

        dist_list.append(distance)
        alpha_diff_list.append(p_a - m_a)
        beta_diff_list.append(p_b - m_b)
        theta_list.append(theta)

    statistics_path = os.path.join(axis_calc_root, STATISTICS_PREFIX + str(man_idx) + TXT_SUFFIX)
    fileutils.write_statistics(statistics_path, dist_list, alpha_diff_list, beta_diff_list, theta_list)

    print(dist_list)
    print('distance mean: ' + str(sum(dist_list) / len(dist_list)))
    print(alpha_diff_list)
    print(beta_diff_list)
    print(theta_list)
    print('theta mean: ' + str(sum(theta_list) / len(theta_list)))


def gen_axis_data(seq_list):

    for seq_number in seq_list:

        # input_path = get_seg_path(seq_number, anti_root)
        input_path = os.path.join(anti_root, PRED_PREFIX + seq_number + NIFTI_SUFFIX)
        output_axis_path = os.path.join(output_axis_root, OUTPUT_AXIS_PREFIX + seq_number + NIFTI_SUFFIX)
        point, vector = find_long_axis.find_long_axis(input_path, out_path=None,
                                                      out_rotated_path=None, ratio_down=1 / 3,
                                                      ratio_up=2 / 3, ratio_start=0.4)
        gen_path = os.path.join(generated_root, GENERATED_PREFIX + seq_number + TXT_SUFFIX)
        fileutils.write_axis(gen_path, point, vector)


def gen_antialias(seq_list, radius=15, ksize=9):
    for seq_number in seq_list:
        aliased_path = os.path.join(aliased_root, PRED_PREFIX + seq_number + NIFTI_SUFFIX)
        antialiased_path = os.path.join(anti_root, PRED_PREFIX + seq_number + NIFTI_SUFFIX)
        t = time.process_time()
        antialias.antialias(aliased_path, antialiased_path, radius, ksize)
        tt = time.process_time()
        print('finished antialias: ' + seq_number, ', cost:', tt - t, 'seconds')


def gen_excel(patients):
    # write to csv
    f_all = open(os.path.join(axis_calc_root, 'statistic.csv'), 'w')
    heads = ['patient', 'm1_a', 'm2_a', 'm3_a', 'DL_alpha', 'm1_b', 'm2_b',
             'm3_b',
             'DL_beta', 'manual1_alpha1', 'manual2_alpha1', 'manual3_alpha1', 'DL_alpha1', 'manual1_relative_angle',
             'manual2_relative_angle',
             'manual3_relative_angle', 'manual1_distance', 'manual2_distance', 'manual3_distance']

    f_all.write(','.join(heads))
    f_all.write('\n')

    for patient in patients:
        patient = str(patient).zfill(5)
        manual_metrics1 = os.path.join(metrics_root[0], METRICS_PREFIX + patient + TXT_SUFFIX)
        manual_metrics2 = os.path.join(metrics_root[1], METRICS_PREFIX + patient + TXT_SUFFIX)
        manual_metrics3 = os.path.join(metrics_root[2], METRICS_PREFIX + patient + TXT_SUFFIX)

        d1, (p1_a, p1_b), (m1_a, m1_b), theta1 = fileutils.read_metrics(manual_metrics1)
        d2, (p2_a, p2_b), (m2_a, m2_b), theta2 = fileutils.read_metrics(manual_metrics2)
        d3, (p3_a, p3_b), (m3_a, m3_b), theta3 = fileutils.read_metrics(manual_metrics3)

        assert p1_a == p2_a
        assert p1_b == p2_b
        assert p1_a == p3_a
        assert p1_b == p3_b

        DL_alpha = p1_a
        DL_beta = p1_b

        manual1_alpha1 = str(abs(90 - abs(float(m1_a))))
        manual2_alpha1 = str(abs(90 - abs(float(m2_a))))
        manual3_alpha1 = str(abs(90 - abs(float(m3_a))))
        DL_alpha1 = str(abs(90 - abs(float(DL_alpha))))

        write_list = list(map(str, [patient, m1_a, m2_a, m3_a, DL_alpha, m1_b, m2_b,
                      m3_b, \
                      DL_beta, manual1_alpha1, manual2_alpha1, manual3_alpha1, DL_alpha1, theta1,
                      theta2, \
                      theta3, d1, d2, d3]))
        f_all.write(','.join(write_list) + '\n')
    f_all.close()


def gen_merged(seq_list, extend_axis_root=None):

    global error_list

    ii = 0

    for seq_number in seq_list:
        machine_seg_path = os.path.join(output_axis_root, OUTPUT_AXIS_PREFIX + seq_number + NIFTI_SUFFIX)
        manual0_seg_path = os.path.join(manual_root[0], seq_number + M0_SUFFIX + NRRD_SUFFIX)
        manual1_seg_path = os.path.join(manual_root[1], seq_number[2:] + M1_SUFFIX + NRRD_SUFFIX)
        manual2_seg_path = os.path.join(manual_root[2], seq_number[2:] + M2_SUFFIX + NRRD_SUFFIX)


        if ii in error_list:
            print('skipped wrong size')
            ii += 1
            continue

        machine_seg = sitk.GetArrayFromImage(sitk.ReadImage(machine_seg_path))
        manual0_seg = sitk.GetArrayFromImage(sitk.ReadImage(manual0_seg_path))
        manual1_seg = sitk.GetArrayFromImage(sitk.ReadImage(manual1_seg_path))
        manual2_seg = sitk.GetArrayFromImage(sitk.ReadImage(manual2_seg_path))

        for i in range(machine_seg.shape[0]):
            for j in range(machine_seg.shape[1]):
                for k in range(machine_seg.shape[2]):
                    if manual0_seg[i, j, k]:
                        machine_seg[i, j, k] = 3.0
                    if manual1_seg[i, j, k]:
                        machine_seg[i, j, k] = 4.0
                    if manual2_seg[i, j, k]:
                        machine_seg[i, j, k] = 5.0

        if extend_axis_root is not None:
            extend_axis_path = os.path.join(extend_axis_root, GENERATED_PREFIX + seq_number + TXT_SUFFIX)
            point, vector = fileutils.read_axis(extend_axis_path)
            find_long_axis.draw_axis(machine_seg, point, vector, 6.0)

        merged_path = os.path.join(merged_root, MERGED_PREFIX + seq_number + NIFTI_SUFFIX)
        sitk.WriteImage(sitk.GetImageFromArray(machine_seg), merged_path)
        print('finished merge: ' + seq_number)
        ii += 1


def rotate_image(image_path, axis_path, rotated_path, center=None, method=0, scaler=None):
    image = sitk.ReadImage(image_path)
    # print_image_info(image)
    ori_size = image.GetSize()

    # image.SetOrigin((0.0, 0.0, 0.0))
    # image.SetSpacing((1.0, 1.0, 1.0))
    # image.SetDirection((1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0))

    point, vec = fileutils.read_axis(axis_path)

    # resize image to get enough length of Z
    if image.GetSize()[2] < image.GetSize()[1]:

        image_array = sitk.GetArrayFromImage(image)
        padding = (image_array.shape[1] - image_array.shape[0]) // 2
        padding_array = np.array([np.zeros((image_array.shape[1], image_array.shape[2])) for i in range(padding)])
        image_array = np.concatenate([padding_array, image_array], axis=0)
        image_array = np.concatenate([image_array, padding_array], axis=0)
        cube_image = sitk.GetImageFromArray(image_array)

        # cube_image.CopyInformation(image)
        cube_image.SetSpacing(image.GetSpacing())
        cube_image.SetDirection(image.GetDirection())
        cube_image.SetOrigin(np.asarray(image.GetOrigin()) - np.asarray([0.0, 0.0, padding]) * np.asarray(image.GetSpacing()))

        image = cube_image
        point[2] += padding


    # scaler adjustment
    vec = np.asarray(image.GetSpacing()) * vec
    vec /= np.linalg.norm(vec)




    if center == None:
        image_array = sitk.GetArrayFromImage(image)
        if method == 0:

            center_z = round(np.mean(find_long_axis.get_exist_range(image_array)))
            image_array = image_array.swapaxes(0, 1)
            center_y = round(np.mean(find_long_axis.get_exist_range(image_array)))
            image_array = image_array.swapaxes(0, 2)
            center_x = round(np.mean(find_long_axis.get_exist_range(image_array)))
            center_heart = (center_x, center_y, center_z)

        elif method == 1:
            l = [0, 0, 0]
            u = [0, 0, 0]
            l[2], u[2] = find_long_axis.get_exist_range(image_array)
            image_array = image_array.swapaxes(0, 1)
            l[1], u[1] = find_long_axis.get_exist_range(image_array)
            image_array = image_array.swapaxes(0, 2)
            l[0], u[0] = find_long_axis.get_exist_range(image_array)
            center_heart = tuple(geoutils.get_axis_center(point, vec, l, u))
        else:
            print('error: no known rotate method!')

    else:
        center_heart = center

    center_origin = (
        image.GetSize()[0] // 2,
        image.GetSize()[1] // 2,
        image.GetSize()[2] // 2
    )

    # trans_translation = sitk.Euler3DTransform()
    x = np.asarray(center_heart) - np.asarray(center_origin)
    x = x.astype(np.float64)
    # trans_translation.SetTranslation(x)
    # image = sitk.Resample(image, trans_translation)

    if method == 0:
        # old-fasion transform
        theta_y, theta_z = geoutils.get_rotation_angle(vec)
        trans_rotate = sitk.Euler3DTransform()
        trans_rotate.SetCenter(np.asarray(center_origin) * np.asarray(image.GetSpacing()) + np.asarray(image.GetOrigin()))
        trans_rotate.SetRotation(0.0, 0.0, theta_z)
        trans_rotate.SetTranslation(np.asarray(x) * np.asarray(image.GetSpacing()))
        out = sitk.Resample(image, trans_rotate, interpolator=sitk.sitkNearestNeighbor)
        trans_rotate.SetRotation(0.0, theta_y, 0.0)
        trans_rotate.SetTranslation((0.0, 0.0, 0.0))
        out = sitk.Resample(out, trans_rotate, interpolator=sitk.sitkNearestNeighbor)

        # old-fasion adjustment
        trans_rotate = sitk.Euler3DTransform()
        trans_rotate.SetCenter(np.asarray(center_origin) * np.asarray(image.GetSpacing()) + np.asarray(image.GetOrigin()))
        trans_rotate.SetRotation(0.0, 0.0, -1 * math.pi / 2)
        out = sitk.Resample(out, trans_rotate, interpolator=sitk.sitkNearestNeighbor)

# =========================================================
    elif method == 1:
        # rodrigues-fasion transform
        trans = sitk.Euler3DTransform()

        # actual calculating logic
        # R = geoutils.get_rodrigues_rotation_matrix(vec, np.asarray([0., 0., 1.]))
        # R = np.linalg.inv(R)

        # changed to save hash power and accuracy
        R = geoutils.get_rodrigues_rotation_matrix(np.asarray([0., 0., 1.]), vec)

        trans.SetMatrix(R.flatten().tolist())
        trans.SetTranslation(np.asarray(x) * np.asarray(image.GetSpacing()))
        trans.SetCenter(np.asarray(center_origin) * np.asarray(image.GetSpacing()) + np.asarray(image.GetOrigin()))
        out = sitk.Resample(image, trans, interpolator=sitk.sitkNearestNeighbor)

        # adjustment
        trans_rotate = sitk.Euler3DTransform()
        trans_rotate.SetCenter(np.asarray(center_origin) * np.asarray(image.GetSpacing()) + np.asarray(image.GetOrigin()))
        trans_rotate.SetRotation(0.0, 0.0, np.arctan(0.5))
        out = sitk.Resample(out, trans_rotate, interpolator=sitk.sitkNearestNeighbor)
# ===============================================

    # out.SetOrigin(ori_origin)
    # out.SetSpacing(ori_spacing)
    # out.SetDirection(ori_direction)

    sitk.WriteImage(out, rotated_path)
    # print_image_info(out)
    return center_heart, ori_size


def get_seg_path(seq_number):
    if __is_pred:
        return os.path.join(recovered_root, PRED_PREFIX + seq_number + NIFTI_SUFFIX)
    else:
        seg_path = os.path.join(image_root, IMAGEDIR_PREFIX + seq_number, SEG_NAME)
        if os.path.isfile(seg_path):
            return seg_path
        return os.path.join(gt_root, SEGMENTATION_PREFIX + seq_number + NIFTI_SUFFIX)


def gen_rotated_shortAxis(seq_list):

    error_list = []

    for seq_number in seq_list:
        image_path = os.path.join(image_root, IMAGEDIR_PREFIX + seq_number, IMAGE_NAME)
        segment_path = os.path.join(anti_root, PRED_PREFIX + seq_number + NIFTI_SUFFIX)
        axis_path = os.path.join(generated_root, GENERATED_PREFIX + seq_number + TXT_SUFFIX)
        rotated_folder = os.path.join(shortAxis_root, IMAGEDIR_PREFIX + seq_number)
        if not os.path.isdir(rotated_folder):
            os.mkdir(rotated_folder)
        rotated_image_path = os.path.join(rotated_folder, IMAGE_NAME)
        rotated_seg_path = os.path.join(rotated_folder, SEG_NAME)

        if not os.path.exists(image_path):
            error_list.append(seq_number)
            continue

        center, size = rotate_image(segment_path, axis_path, rotated_seg_path)
        
        fileutils.write_tempinfo(os.path.join(tempinfo_root, TEMPINFO_PREFIX + seq_number + TXT_SUFFIX), center, size)

        rotate_image(image_path, axis_path, rotated_image_path, center=center)

        print('finished rotate to short axis: ' + seq_number)
    print('error:', error_list)

def rotate_back(source_path, axis_path, tempinfo_path, target_path, method=0):
    source_img = sitk.ReadImage(source_path)
    center_heart, ori_size = fileutils.read_tempinfo(tempinfo_path)

    center_origin = (
        source_img.GetSize()[0] // 2,
        source_img.GetSize()[1] // 2,
        source_img.GetSize()[2] // 2
    )

    _, vec = fileutils.read_axis(axis_path)

    vec = np.asarray(source_img.GetSpacing()) * vec
    vec /= np.linalg.norm(vec)

    target_image = source_img

    if method == 0:
        # old-fasion transform
        theta_y, theta_z = geoutils.get_rotation_angle(vec)
        trans = sitk.Euler3DTransform()
        trans.SetRotation(0.0, 0.0, math.pi / 2)
        trans.SetCenter(np.asarray(center_origin) * np.asarray(source_img.GetSpacing()) + np.asarray(source_img.GetOrigin()))
        target_image = sitk.Resample(target_image, trans, interpolator=sitk.sitkNearestNeighbor)
        trans.SetRotation(0.0, -theta_y, 0.0)
        target_image = sitk.Resample(target_image, trans, interpolator=sitk.sitkNearestNeighbor)
        trans.SetRotation(0.0, 0.0, -theta_z)
        target_image = sitk.Resample(target_image, trans, interpolator=sitk.sitkNearestNeighbor)
        trans.SetRotation(0.0, 0.0, 0.0)
        x = np.asarray(center_origin) - np.asarray(center_heart)
        x = x.astype(np.float64)
        trans.SetTranslation(np.asarray(x) * np.asarray(source_img.GetSpacing()))
        target_image = sitk.Resample(target_image, trans, interpolator=sitk.sitkNearestNeighbor)

# ======================================
    elif method == 1:
        # rodrigues-fasion transform
        # adjustment
        trans = sitk.Euler3DTransform()
        trans.SetCenter(np.asarray(center_origin) * np.asarray(source_img.GetSpacing()) + np.asarray(source_img.GetOrigin()))
        trans.SetRotation(0.0, 0.0, - np.arctan(0.5))
        target_image = sitk.Resample(target_image, trans, interpolator=sitk.sitkNearestNeighbor)

        R = geoutils.get_rodrigues_rotation_matrix(vec, np.asarray([0., 0., 1.]))
        trans.SetMatrix(R.flatten().tolist())
        target_image = sitk.Resample(target_image, trans, interpolator=sitk.sitkNearestNeighbor)

        trans.SetRotation(0.0, 0.0, 0.0)
        x = np.asarray(center_origin) - np.asarray(center_heart)
        x = x.astype(np.float64)
        trans.SetTranslation(np.asarray(x) * np.asarray(source_img.GetSpacing()))
        target_image = sitk.Resample(target_image, trans, interpolator=sitk.sitkNearestNeighbor)

# ======================================

    else:
        print('error: no known rotate method')

    target_array = sitk.GetArrayFromImage(target_image)
    padding = int((target_image.GetSize()[2] - ori_size[2]) // 2)
    if padding != 0:
        target_array = target_array[padding:-padding]
    target_image = sitk.GetImageFromArray(target_array)
    target_image.SetOrigin(np.asarray(source_img.GetOrigin()) + np.asarray([0.0, 0.0, padding]) * source_img.GetSpacing())
    target_image.SetDirection(source_img.GetDirection())
    target_image.SetSpacing(source_img.GetSpacing())

    sitk.WriteImage(target_image, target_path)


def gen_rotate_seg_back(seq_list):
    '''
    need the seg has the correct scaler!

    :param seq_list:
    :return:
    '''

    for seq_number in seq_list:

        source_path = os.path.join(shortAxisSeg_root, SEGMENTATION_PREFIX + seq_number + NIFTI_SUFFIX)
        axis_path = os.path.join(generated_root, GENERATED_PREFIX + seq_number + TXT_SUFFIX)
        tempinfo_path = os.path.join(tempinfo_root, TEMPINFO_PREFIX + seq_number + TXT_SUFFIX)
        target_path = os.path.join(rotateBack_root, SEGMENTATION_PREFIX + seq_number + NIFTI_SUFFIX)
        rotate_back(source_path, axis_path, tempinfo_path, target_path)
        print('finished rotate back:', seq_number)



error_list = []
def check_size():
    man_files = os.listdir(manual_root[0])
    global error_list
    for idx in range(1, 3):
        target_files = os.listdir(manual_root[idx])
        for i in range(len(man_files)):
            standard = sitk.ReadImage(os.path.join(manual_root[0], man_files[i]))
            target = sitk.ReadImage(os.path.join(manual_root[idx], target_files[i]))
            if standard.GetSize() != target.GetSize():
                print('wrong size, man_idx: ', idx, ', i: ', i + 177)
                error_list.append(i)


def gen_templateinfo(seq_list):
    for seq_number in seq_list:
        image_path = os.path.join(image_root, IMAGEDIR_PREFIX + seq_number, IMAGE_NAME)
        image = sitk.ReadImage(image_path)
        template_path = os.path.join(templateinfo_root, TEMPLATEINFO_PREFIX + seq_number + TXT_SUFFIX)
        fileutils.write_template_info(template_path, image.GetSize(), image.GetSpacing(), image.GetOrigin(), image.GetDirection())
        print('finished writing templateinfo:', seq_number)


# need set the resampled, recovered root as status input
def gen_recovered(seq_list):
    #TODO: read the origin info to recovered the downsampled image
    for seq_number in seq_list:
        img_path = os.path.join(resampled_root, PRED_PREFIX + seq_number + NIFTI_SUFFIX)
        templateinfo_path = os.path.join(templateinfo_root, TEMPLATEINFO_PREFIX + seq_number + TXT_SUFFIX)
        src_img = sitk.ReadImage(img_path)
        src_img = sitk.Cast(src_img, sitk.sitkFloat32)
        
        tpl_size, tpl_spacing, tpl_origin, tpl_direction = fileutils.read_template_info(templateinfo_path)

        #TODO: fake the spacing and origin of template, may have problem?
        # tpl_spacing = (1.0, 1.0, 1.0)
        # tpl_origin = (0.0, 0.0, 0.0)

        src_spacing = np.asarray(tpl_spacing) * np.asarray(tpl_size) / np.asarray(src_img.GetSize())
        
        # adjust the direction
        # src_spacing *= np.asarray(tpl_direction) / np.asarray(src_img.GetDirection())
        # src_spacing *= np.asarray([-1, -1, 1])
         

        src_img.SetSpacing(src_spacing)
        src_img.SetOrigin(tpl_origin)
        src_img.SetDirection(tpl_direction)

        print(tpl_size)
        print(tpl_spacing)
        print(tpl_origin)

        target_img = sitk.Resample(src_img, tpl_size, sitk.Euler3DTransform(), sitk.sitkNearestNeighbor, tpl_origin, tpl_spacing, tpl_direction)

        target_path = os.path.join(recovered_root, PRED_PREFIX + seq_number + NIFTI_SUFFIX)
        sitk.WriteImage(target_img, target_path)
        print('recovered', seq_number)


def gen_correct_dir(seq_list):
    #TODO:
    for seq_number in seq_list:
        in_path = os.path.join(resampled_root, PRED_PREFIX + seq_number + NIFTI_SUFFIX)
        image = sitk.ReadImage(in_path)
        templateinfo_path = os.path.join(templateinfo_root, TEMPLATEINFO_PREFIX + seq_number + TXT_SUFFIX)
        _1, _2, _3, direction = fileutils.read_template_info(templateinfo_path)
        image.SetDirection(direction)
        sitk.WriteImage(image, in_path)
        print('correct the direction,', seq_number)


def gen_post_processed(seq_list):
    #TODO:
    for seq_number in seq_list:
        in_path = get_seg_path(seq_number)
        out_path = os.path.join(recovered_root, PRED_PREFIX + seq_number + NIFTI_SUFFIX)
        # put the post processed seg in recovered_root, even though the seg is from image_root
        post_process.post_process(in_path, out_path)
        print('finished post process for', seq_number)


def set_scaler(target_img, template_img):
    target_img.SetOrigin(template_img.GetOrigin())
    target_img.SetSpacing(template_img.GetSpacing())
    target_img.SetDirection(template_img.GetDirection())


def batch_set_scaler(seq_list):
    for seq_number in seq_list:
        image_path = os.path.join(image_root, IMAGEDIR_PREFIX + seq_number, IMAGE_NAME)
        seg_path = os.path.join(anti_root, PRED_PREFIX + seq_number + NIFTI_SUFFIX)
        image = sitk.ReadImage(image_path)
        seg = sitk.ReadImage(seg_path)
        set_scaler(seg, image)
        sitk.WriteImage(seg, seg_path)
        print('finished set scaler:', seq_number)


def upsampling_procedure(seq_list):
    if __is_pred:
        gen_templateinfo(seq_list)
        gen_correct_dir(seq_list)
        gen_recovered(seq_list)
        pass
    gen_post_processed(seq_list)
    gen_antialias(seq_list)


def axis_calc_procedure(seq_list):
    gen_axis_data(seq_list)
    pass


def axis_evaluate_procedure(seq_list):
    for i in range(3):
        calc_metrics(seq_list, i)
        calc_statistics(seq_list, i)
    # gen_excel(seq_list)


def shortaxis_rotate_procedure(seq_list):
    # let segmentation have the same scaler with the image
    batch_set_scaler(seq_list)
    gen_rotated_shortAxis(seq_list)


def rotate_back_procedure(seq_list):
    gen_rotate_seg_back(seq_list)


def task(seq_list):
    np.set_printoptions(precision=16)

    # upsampling_procedure(seq_list)

    # axis_calc_procedure(seq_list)

    shortaxis_rotate_procedure(seq_list)

    # rotate_back_procedure(seq_list)


    # check_size()
    # print(error_list)
    # gen_merged(seq_list)


if __name__ == '__main__':
    train_list = [144, 146, 147, 149, 152, 154, 155, 158, 159, 160, 161, 162, 166, 167, 168, 170, 171, 172, 173, 176]
    train_list += list(range(200, 258))

    ori_test_list = list(range(177, 197))
    ori_test_list.remove(186)

    extra_test_list = list(range(260, 345))
    extra_test_list += list(range(350, 442))


    seq_list = [str(i).zfill(5) for i in extra_test_list]
    seq_list = ['00441']

    task(seq_list)

    # processor_cnt = max(mp.cpu_count() - 2, 1)
    # 
    # processor_size = (len(seq_list) + processor_cnt - 1) // processor_cnt
    # args_list = []
    # lptr = 0
    # while lptr < len(seq_list):
    #     args_list.append(seq_list[lptr:lptr + processor_size])
    #     lptr += processor_size
    # print(len(args_list))
    # print(processor_size)
    # pool = mp.Pool(processor_cnt)
    # pool.map(task, args_list)

    # pool.close()
    # pool.join()

    # axis_evaluate_procedure(seq_list)
    # gen_excel(seq_list)
